package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
	"time"

	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"strconv"

	"golang.org/x/crypto/chacha20"
	"golang.org/x/crypto/chacha20poly1305"

	"github.com/fsnotify/fsnotify"
)

// go run main.go sync.go crypto.go
const (
	ModeMaster              = "Host"
	ModeSlave               = "Slave"
	HeartbeatInterval       = 5 * time.Second
	InitialTimeout          = 15 * time.Second
	SlaveTimeout            = 15 * time.Second
	HeartbeatMessage        = "HEARTBEAT"
	HeartbeatResponse       = "OK"
	HeartbeatValidityPeriod = 10 * time.Second // 心跳包有效期为10秒

)

type Config struct {
	IPAddress string `json:"ip_address"`
	HPort     int    `json:"h_port"`
	PeerIP    string `json:"peer_ip"`
	SyncPort  int    `json:"sync_port"`
	SecretKey string `json:"secret_key"`
	Crypto    string `json:"crypto"` // 新增字段，用来指定加解密算法
	Enable    bool   `json:"enable"`
}

type metric struct {
	heartbeatsSent     int64     // 心跳包发送次数
	heartbeatsReceived int64     // 心跳包接收次数
	syncAttempts       int64     // 同步尝试次数
	syncSuccess        int64     // 同步成功次数
	syncFailed         int64     // 同步失败次数
	lastSyncTime       time.Time // 最后同步时间
}

var (
	config            Config
	modeLock          sync.Mutex
	currentMode       string
	metrics           = metric{} // 用于存储指标
	HeartbeatValidity = false    // 心跳包有效性标志
)

func main() {
	if err := loadConfig(); err != nil {
		log.Fatal("加载配置文件失败:", err)
	}
	fmt.Println("启动从机模式，监听 UDP 心跳...")
	go listenUDPHeartbeat()

	// 程序启动直接进入从机模式
	startSlaveModeSwitch()
	// 启动从机同步模式
	go receiveFileFromHost()

	select {}
}

func loadConfig() error {
	basePath, err := os.Getwd()
	if err != nil {
		return err
	}
	log.Printf("当前工作目录: %s", basePath)

	configPath := filepath.Join(basePath, "config", "mangos_config.json")
	file, err := os.Open(configPath)
	if err != nil {
		return err
	}
	defer file.Close()

	decoder := json.NewDecoder(file)
	return decoder.Decode(&config)
}

func setMode(mode string) {
	modeLock.Lock()
	defer modeLock.Unlock()
	if currentMode != mode {
		currentMode = mode
		fmt.Printf("切换到 %s 模式\n", mode)
	}
}

func startHostMode() {
	setMode(ModeMaster)
	go sendUDPHeartbeatToPeer() // 向 peer_ip 和 h_port 发送心跳包
	go syncFiles()
}

func sendUDPHeartbeatToPeer() {
	// 设置 UDP 地址和端口
	addr := net.JoinHostPort(config.PeerIP, fmt.Sprintf("%d", config.HPort))
	conn, err := net.Dial("udp", addr)
	if err != nil {
		log.Fatal("无法连接到目标:", err)
	}
	defer conn.Close()

	for {
		heartbeatMessage := generateHeartbeatMessage()
		if heartbeatMessage == "" {
			continue
		}

		_, err = conn.Write([]byte(heartbeatMessage))
		if err != nil {
			log.Printf("发送心跳包失败: %v", err)
		}

		log.Printf("发送心跳包：%s", heartbeatMessage)
		atomic.AddInt64(&metrics.heartbeatsSent, 1) // 更新心跳包发送次数
		time.Sleep(HeartbeatInterval)
	}
}

func listenUDPHeartbeat() {
	addr := fmt.Sprintf(":%d", config.HPort) // 使用配置中的 HPort
	var conn net.PacketConn
	var err error

	// 尝试绑定端口，最多重试5次
	for i := 0; i < 5; i++ {
		conn, err = net.ListenPacket("udp", addr)
		if err == nil {
			break
		}
		log.Printf("无法监听 UDP 端口 %d: %v. 正在重试 (%d/5)...", config.HPort, err, i+1)
		time.Sleep(2 * time.Second)
	}

	if err != nil {
		log.Fatalf("重试 5 次后仍无法监听 UDP 端口 %d: %v", config.HPort, err)
	}
	defer conn.Close()

	buffer := make([]byte, 1024)
	for {
		n, addr, err := conn.ReadFrom(buffer)
		if err != nil {
			log.Println("读取 UDP 消息失败:", err)
			continue
		}

		receivedMessage := string(buffer[:n])
		// 打印心跳包来源
		log.Printf("接收到 UDP 心跳: %s 来自远程地址 %s", receivedMessage, addr)

		if validateHeartbeatMessage(receivedMessage) {
			atomic.AddInt64(&metrics.heartbeatsReceived, 1) // 更新心跳包接收次数
			HeartbeatValidity = true
		} else {
			log.Println("接收到的心跳包无效:", receivedMessage)
			HeartbeatValidity = false
		}
	}
}

func startSlaveModeSwitch() {
	setMode(ModeSlave)

	// 启动从机模式的心跳检测逻辑
	go func() {
		timer := time.NewTimer(SlaveTimeout)
		defer timer.Stop()

		for {
			select {
			case <-timer.C:
				log.Println("未收到心跳包，切换到主机模式...")
				startHostMode()
				return
			default:
				// 检查是否收到心跳包
				if HeartbeatValidity {
					//log.Println("收到心跳包，重置计时器...")
					timer.Reset(SlaveTimeout) // 重置计时器
					HeartbeatValidity = false // 重置心跳包有效性标志
				}
				time.Sleep(2 * time.Second) // 检查间隔
			}
		}
	}()
}

// 规范化密钥
func normalizeSecretKey(secretKey string) string {
	// 确保密钥长度为 16、24 或 32 字节
	switch len(secretKey) {
	case 16, 24, 32:
		return secretKey
	default:
		if len(secretKey) > 32 {
			return secretKey[:32] // 截断到 32 字节
		}
		// 填充到最近的有效长度
		padding := 16
		if len(secretKey) > 16 {
			padding = 24
		}
		return secretKey + string(make([]byte, padding-len(secretKey)))
	}
}

// 加密心跳包（AES-GCM）
func encryptHeartbeatMessageAES(secretKey, message string) (string, error) {
	normalizedKey := normalizeSecretKey(secretKey) // 规范化密钥
	block, err := aes.NewCipher([]byte(normalizedKey))
	if err != nil {
		return "", err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	ciphertext := gcm.Seal(nonce, nonce, []byte(message), nil)
	return hex.EncodeToString(ciphertext), nil
}

// 解密心跳包（AES-GCM）
func decryptHeartbeatMessageAES(secretKey, encryptedMessage string) (string, error) {
	normalizedKey := normalizeSecretKey(secretKey) // 规范化密钥
	block, err := aes.NewCipher([]byte(normalizedKey))
	if err != nil {
		return "", err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	ciphertext, err := hex.DecodeString(encryptedMessage)
	if err != nil {
		return "", err
	}

	if len(ciphertext) < gcm.NonceSize() {
		return "", errors.New("ciphertext too short")
	}

	nonce := ciphertext[:gcm.NonceSize()]
	ciphertext = ciphertext[gcm.NonceSize():]

	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return "", err
	}

	return string(plaintext), nil
}

// 加密心跳包（ChaCha20）
func encryptHeartbeatMessageChaCha20(secretKey, message string) (string, error) {
	// ChaCha20 需要 256 位的密钥
	key := []byte(secretKey)
	if len(key) < 32 {
		key = append(key, make([]byte, 32-len(key))...)
	}

	var nonce [chacha20.NonceSize]byte
	_, err := io.ReadFull(rand.Reader, nonce[:])
	if err != nil {
		return "", err
	}

	c, err := chacha20.NewUnauthenticatedCipher(key, nonce[:])
	if err != nil {
		return "", err
	}

	plaintext := []byte(message)
	ciphertext := make([]byte, len(plaintext))
	c.XORKeyStream(ciphertext, plaintext)

	// 将加密的密文和 nonce 组合起来并转为 hex 编码
	return hex.EncodeToString(append(nonce[:], ciphertext...)), nil
}

// 解密心跳包（ChaCha20）
func decryptHeartbeatMessageChaCha20(secretKey, encryptedMessage string) (string, error) {
	// ChaCha20 需要 256 位的密钥
	key := []byte(secretKey)
	if len(key) < 32 {
		key = append(key, make([]byte, 32-len(key))...)
	}

	ciphertext, err := hex.DecodeString(encryptedMessage)
	if err != nil {
		return "", err
	}

	if len(ciphertext) < chacha20.NonceSize {
		return "", errors.New("ciphertext too short")
	}

	// 获取 nonce 和密文部分
	nonce := ciphertext[:chacha20.NonceSize]
	ciphertext = ciphertext[chacha20.NonceSize:]

	c, err := chacha20.NewUnauthenticatedCipher(key, nonce)
	if err != nil {
		return "", err
	}

	plaintext := make([]byte, len(ciphertext))
	c.XORKeyStream(plaintext, ciphertext)

	return string(plaintext), nil
}

// 加密心跳包（ChaCha20-Poly1305）
func encryptHeartbeatMessageChaCha20Poly1305(secretKey, message string) (string, error) {
	// ChaCha20-Poly1305 需要 256 位的密钥
	key := []byte(secretKey)
	if len(key) < 32 {
		key = append(key, make([]byte, 32-len(key))...)
	}

	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return "", err
	}

	nonce := make([]byte, aead.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	ciphertext := aead.Seal(nonce, nonce, []byte(message), nil)
	return hex.EncodeToString(ciphertext), nil
}

// 解密心跳包（ChaCha20-Poly1305）
func decryptHeartbeatMessageChaCha20Poly1305(secretKey, encryptedMessage string) (string, error) {
	// ChaCha20-Poly1305 需要 256 位的密钥
	key := []byte(secretKey)
	if len(key) < 32 {
		key = append(key, make([]byte, 32-len(key))...)
	}

	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return "", err
	}

	ciphertext, err := hex.DecodeString(encryptedMessage)
	if err != nil {
		return "", err
	}

	if len(ciphertext) < aead.NonceSize() {
		return "", errors.New("ciphertext too short")
	}

	nonce, ciphertext := ciphertext[:aead.NonceSize()], ciphertext[aead.NonceSize():]

	plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return "", err
	}

	return string(plaintext), nil
}

// 根据配置选择加解密算法
func encryptHeartbeatMessage(message string) (string, error) {
	switch config.Crypto {
	case "AES-GCM":
		return encryptHeartbeatMessageAES(config.SecretKey, message)
	case "ChaCha20":
		return encryptHeartbeatMessageChaCha20(config.SecretKey, message)
	case "ChaCha20-Poly1305":
		return encryptHeartbeatMessageChaCha20Poly1305(config.SecretKey, message)
	default:
		return "", fmt.Errorf("不支持的加密算法: %s", config.Crypto)
	}
}

func decryptHeartbeatMessage(encryptedMessage string) (string, error) {
	switch config.Crypto {
	case "AES-GCM":
		return decryptHeartbeatMessageAES(config.SecretKey, encryptedMessage)
	case "ChaCha20":
		return decryptHeartbeatMessageChaCha20(config.SecretKey, encryptedMessage)
	case "ChaCha20-Poly1305":
		return decryptHeartbeatMessageChaCha20Poly1305(config.SecretKey, encryptedMessage)
	default:
		return "", fmt.Errorf("不支持的解密算法: %s", config.Crypto)
	}
}

// 生成心跳包
func generateHeartbeatMessage() string {
	timestamp := time.Now().Unix() // 当前时间戳
	message := fmt.Sprintf("%d", timestamp)
	encryptedMessage, err := encryptHeartbeatMessage(message)
	if err != nil {
		log.Println("生成心跳包失败:", err)
		return ""
	}
	return encryptedMessage
}

// 验证心跳包
func validateHeartbeatMessage(encryptedMessage string) bool {
	decryptedMessage, err := decryptHeartbeatMessage(encryptedMessage)
	if err != nil {
		log.Println("解密心跳包失败:", err)
		return false
	}

	timestamp, err := strconv.ParseInt(decryptedMessage, 10, 64)
	if err != nil {
		log.Println("解析心跳包时间戳失败:", err)
		return false
	}
	messageTime := time.Unix(timestamp, 0)
	log.Printf("解密心跳包成功，时间戳: %v 本地时间 %v", timestamp, time.Now().Unix())

	if time.Since(messageTime) > HeartbeatValidityPeriod {
		log.Println("心跳包已过期:", messageTime)
		return false
	}

	return true
}

// 定义 syncFiles 函数
func syncFiles() {
	// 创建 fsnotify 实例
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatalf("无法创建文件监视器: %v", err)
	}
	defer watcher.Close()

	// 监视 ./sync/ 目录下的所有文件变化
	err = watcher.Add("./config/test")
	if err != nil {
		log.Fatalf("无法监视 ./test 目录: %v", err)
	}

	// 创建 UDP 连接，用于发送文件数据
	addr := net.JoinHostPort(config.PeerIP, fmt.Sprintf("%d", config.SyncPort))
	conn, err := net.Dial("udp", addr)
	if err != nil {
		log.Fatalf("无法连接到目标: %v", err)
	}
	defer conn.Close()

	log.Println("开始监听 ./config/test 目录中的文件变化...")

	// 无限循环监听文件变化
	for {
		select {
		case event := <-watcher.Events:
			// 处理文件变化事件
			if event.Op&fsnotify.Create == fsnotify.Create {
				log.Printf("文件创建: %s", event.Name)
				go sendFileToPeer(event.Name, conn) // 文件创建时，发送文件
			}
			if event.Op&fsnotify.Write == fsnotify.Write {
				log.Printf("文件修改: %s", event.Name)
				go sendFileToPeer(event.Name, conn) // 文件修改时，重新发送文件
			}
			if event.Op&fsnotify.Remove == fsnotify.Remove {
				log.Printf("文件删除: %s", event.Name)
				// 可以选择处理删除文件的逻辑，例如发送删除消息给 peer
			}

		case err := <-watcher.Errors:
			log.Printf("文件监视器错误: %v", err)
		}
	}
}

// 发送文件到 peer
func sendFileToPeer(filePath string, conn net.Conn) {
	// 打开文件
	file, err := os.Open(filePath)
	if err != nil {
		log.Printf("无法打开文件 %s: %v", filePath, err)
		return
	}
	defer file.Close()

	// 读取文件内容
	fileInfo, err := file.Stat()
	if err != nil {
		log.Printf("无法获取文件信息 %s: %v", filePath, err)
		return
	}

	fileContent := make([]byte, fileInfo.Size())
	_, err = file.Read(fileContent)
	if err != nil {
		log.Printf("无法读取文件内容 %s: %v", filePath, err)
		return
	}

	// 将文件名和内容组合成消息
	message := fmt.Sprintf("%s|%s", fileInfo.Name(), string(fileContent))

	// 将文件内容发送到 peer
	_, err = conn.Write([]byte(message))
	if err != nil {
		log.Printf("无法发送文件到 %s: %v", config.PeerIP, err)
		return
	}

	log.Printf("文件 %s 内容已发送", filePath)
}

func receiveFileFromHost() {

	if !HeartbeatValidity {
		log.Println("心跳包无效，暂停接收文件数据...")
		return
	}
	addr := fmt.Sprintf(":%d", config.SyncPort) // 使用配置中的 SyncPort
	conn, err := net.ListenPacket("udp", addr)
	if err != nil {
		log.Fatalf("无法监听 UDP 端口 %d: %v", config.SyncPort, err)
	}
	defer conn.Close()

	log.Printf("开始监听来自主机的文件传输，端口: %d", config.SyncPort)

	buffer := make([]byte, 4096) // 缓冲区大小
	for {
		n, remoteAddr, err := conn.ReadFrom(buffer)
		if err != nil {
			log.Printf("接收文件数据失败: %v", err)
			continue
		}

		log.Printf("接收到来自 %s 的文件数据，大小: %d 字节", remoteAddr, n)

		// 分离文件名和内容
		data := string(buffer[:n])
		parts := splitFileMessage(data)
		if len(parts) != 2 {
			log.Println("接收到的文件数据格式无效")
			continue
		}

		fileName, fileContent := parts[0], parts[1]

		// 保存接收到的文件数据到本地
		err = saveReceivedFile(fileName, []byte(fileContent))
		if err != nil {
			log.Printf("保存接收到的文件失败: %v", err)
		} else {
			log.Printf("文件 %s 保存成功", fileName)
		}
	}
}

func saveReceivedFile(fileName string, data []byte) error {
	// 定义保存文件的路径
	savePath := filepath.Join("./config/test", fileName)
	err := os.MkdirAll(filepath.Dir(savePath), os.ModePerm)
	if err != nil {
		return fmt.Errorf("无法创建目录: %v", err)
	}

	file, err := os.Create(savePath)
	if err != nil {
		return fmt.Errorf("无法创建文件: %v", err)
	}
	defer file.Close()

	_, err = file.Write(data)
	if err != nil {
		return fmt.Errorf("写入文件失败: %v", err)
	}

	return nil
}

func splitFileMessage(data string) []string {
	// 分离文件名和内容，假设使用 "|" 作为分隔符
	return filepath.SplitList(data)
}
